"""
模型训练主程序
功能：执行训练循环、保存最佳模型、监控指标
"""
import torch
import torch.optim as optim
from data.loader import get_dataloader, ComponentDataset
from models.multimodel_net import MultiModalNet
from models.physics_loss import PhysicsLoss
import os

def main():
    # 配置参数
    config = {
        'batch_size': 64,
        'lr': 1e-4,
        'epochs': 100,
        'model_save_path': 'outputs/best_model.pth'
    }

    # 创建保存模型的目录
    save_dir = os.path.dirname(config['model_save_path'])
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    # 初始化
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = MultiModalNet().to(device)
    optimizer = optim.AdamW(model.parameters(), lr=config['lr'])
    physics_criterion = PhysicsLoss()

    # 数据加载
    dataset = ComponentDataset('data/component_data.csv')
    dataset.save_preprocessors('outputs')
    train_loader = get_dataloader('data/component_data.csv', config['batch_size'])

    # 训练循环
    best_loss = float('inf')
    for epoch in range(config['epochs']):
        model.train()
        total_loss = 0.0
        for batch in train_loader:
            wave = batch['wave'].to(device)
            struct = batch['struct'].to(device)
            labels = batch['label'].to(device)

            # 前向传播
            pred = model(wave, struct)

            # 计算损失
            mse_loss = torch.nn.functional.mse_loss(pred.flatten(), labels)
            physics_loss = physics_criterion(pred, struct, strain=0.001)  # 假设应变0.1%
            loss = 0.7 * mse_loss + 0.3 * physics_loss

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        # 保存最佳模型
        avg_loss = total_loss / len(train_loader)
        if avg_loss < best_loss:
            torch.save(model.state_dict(), config['model_save_path'])
            best_loss = avg_loss

        print(f"Epoch {epoch + 1}/{config['epochs']} | Loss: {avg_loss:.4f}")


if __name__ == "__main__":
    main()